require "lfs"

PSO={
	iterations = 50,
	inertia = 1.0,
	correction_factor = 2.0,
	swarm_size = 30,
	max_val = 1e20,
	num_iter = 0,
	par = { {name="a", pmax=3000,pmin=1}};
	cache_data = {},
	swarm = {},
	results_path = "iters.txt",
	cache_path = "cache.txt",
	init_type = "random",
};

function PSO:new (o)
	o = o or {};
	self.__index = self;
	setmetatable(o,self)
	return o;
end


function PSO:init( )
	math.randomseed(0)	
	self:init_swarm_array();
	self:read_cache();
end



function random_pos( d )
     return d.pmin+(d.pmax-d.pmin)*math.random()
end


function string.split(s, delimiter)
	result = {};
	for match in (s..delimiter):gmatch("(.-)"..delimiter) do
	    table.insert(result, match);
	end
	return result;
end

function print_table_val( self )
     for k,v in pairs (self) do
	  print( #self, k, v, tonumber(v) )
     end
end
function print_table( self )
     for k,v in pairs (self) do
	  print( #self, k, v )
     end
end
function print_table_table( self )
     for k,v in pairs (self) do
          local str;
	  for kk,vv in pairs( v ) do 
		str = (str or "")..kk..":"..vv..","
	  end
	  print( str )
     end
end
function table2number( self )
     for k,v in pairs (self) do
--~ 	  print( #self, k, v, tonumber(v) )
	  v = tonumber(v)
     end
      return self;
end

function PSO:print( )
	print( "swarm_size:"..self.swarm_size);
	print( "iterations:"..self.iterations);
	print( "inertia:"..self.inertia);
	print( "correction_factor:"..self.correction_factor);
	print( "max_val:"..self.max_val);
	print( "num_iter:"..self.num_iter);
	print( "swarm arr:"..#self.swarm)
	print( "init_type:"..#self.init_type)
	print( "parameter size:"..#self.par)
	print( "parameters:")
	print_table_table(self.par)
end

function PSO:read_cache( path )
	local file = io.open( self.cache_path,"r");
	if file==nile then return; end
	assert( file, "file:"..self.cache_path.." open fail!");
	for line in file:lines() do
--~ 	     print( line )
	     vals = string.split( line, "\t" ) 
	    vals = table2number( vals )
--~ 	     print_table_val( vals );
	     table.insert(self.cache_data, vals)
        end
	file:close()    
	print( "read line:", #self.cache_data );
	print_table_table( self.cache_data );
end

function PSO:dump_particle( particle )
     local data={};
      data[1] = self.num_iter;
      data[2] = particle.index
      for k,v in pairs(self.par) do
           data[2+k] = particle[k].pos;	
	   data[2+k + #self.par] = particle[k].vel; 
	   data[2+k + 2*(#self.par)] = particle[k].best_pos; 
     end
     return data;
end
------------------------------------------------
--查找缓存是否存在已经计算完成的数据，如果存在，返回data
--如果不存在，返回nil
function PSO:find_cache( particle )
     data = self:dump_particle( particle )
     for k,v in pairs( self.cache_data ) do
           local flag = true;
	   for kd,vk in pairs( data ) do 	       
		if math.abs( vk - (v[kd]) ) > 1e-8 then
		     flag = false;		    
		end
	   end
	   if flag then
--~ 	        print_table( v )
		return v;
	   end
     end
     return nil;
end
------------------------------------------------

------------------------------------------------
--write the calculation data of partical to file
function PSO:write2file( particle )
--~ 	print( self.results_path );
	local file = io.open(self.results_path, "a")
	file:write(self.num_iter,"\t", particle.index)
	for k,v in pairs(self.par) do
             file:write("\t",particle[k].pos)
        end
	for k,v in pairs(self.par) do
             file:write("\t",particle[k].vel)
        end
	for k,v in pairs(self.par) do
             file:write("\t",particle[k].best_pos)
        end
        file:write("\t",particle.value,"\n")
	file:close()	
end
------------------------------------------------



---------------------------------------------------------------------------------
--初始化数组
function PSO:init_swarm_array()	
	self.swarm={};
	for ii=1,self.swarm_size do
	    self.swarm[ii]={}
	    for k,v in pairs(self.par) do
		self.swarm[ii][k] = {}
		if self.init_type == "random" then
			self.swarm[ii][k].pos = random_pos(v);  --初始化位置
		else
			self.swarm[ii][k].pos = v.pos;  --初始化位置
		end
		self.swarm[ii][k].vel = 0;  --初始化速度 initial velocity
		self.swarm[ii][k].best_pos = 0;  --初始化位置	
	    end
	    self.swarm[ii].best_value = self.max_val;  --初始化适用度
	    self.swarm[ii].index = ii;  --粒子编号
	end	
end
---------------------------------------------------------------------------------



function PSO:print_swarm_array(num)
	if type(num) == "number" and num > 1 and num < self.swarm_size then		
		print("index:"..num..":")
		    for k,v in pairs(self.par) do
			--swarm[ii][k].pos = myradome(v);
			print( "  k:"..k.." "..v.name..":" )
			print( "    pos:"..self.swarm[num][k].pos);
		end
	else
		for ii=1,self.swarm_size do
		    print("index:"..ii..":")
		    for k,v in pairs(self.par) do
			--swarm[ii][k].pos = myradome(v);
			print( "  k:"..k.." "..v.name..":" )
			print( "    pos:"..self.swarm[ii][k].pos);
		    end
		end
	end
end


-------------------------------------------------------------------------------------------------
--找出最优粒子
function PSO:get_global_best_value()
    local min=self.swarm[1].best_value;
    local idx = 1;
    for ii=2,self.swarm_size do
        if min >= self.swarm[ii].best_value then
	    min = self.swarm[ii].best_value;
	    idx = ii;
	end
    end
    return self.swarm[idx],min,idx;
end
---------------------------------------------------------------------------------------------------


function PSO:envaluation( partical )
     local val = self.max_val;
     for k,v in pairs(self.par) do
         if partical[k].pos < v.pmin or partical[k].pos > v.pmax then
	     return self.max_val;
	 end	 
     end
     return math.pow((partical[1].pos - 15),2)
end

function PSO:check_boundary( partical )     
     for k,v in pairs(self.par) do
         if partical[k].pos < v.pmin or partical[k].pos > v.pmax then
	     return false;
	 end	 
     end
     return true;
end

function PSO:print_partical( partical )
    local str="";
    for k,v in pairs(self.par) do
        str = str..v.name.."="..tostring(partical[k].best_pos).."  ";
    end
    str = str..partical.best_value
    print(str);
end

function PSO:step()	
	self.num_iter = self.num_iter+1;	 	
	print( "num_iter:",self.num_iter)
	 	
	----------------------------------------------------------------------------------------------------------------------------------------
	 --update  position 
	 for jj = 1,self.swarm_size do
		data = std.vector:new( table.maxn(self.par) );
		 if mpi.world.rank() == 0 then 			
			for k,v in pairs(self.par) do
				self.swarm[jj][k].pos = self.swarm[jj][k].pos + self.swarm[jj][k].vel /1.3; -- update  position
				data[k-1] = self.swarm[jj][k].pos;
			end		
			mpi.world.broadcast(data,0)
		else
			mpi.world.broadcast(data,0);
			if( mpi.world.rank() ~= 0 ) then 			
				for k,v in pairs(self.par) do
					self.swarm[jj][k].pos = data[k-1];
				end				
			end
		end	
	end
	----------------------------------------------------------------------------------------------------------------------------------------
--~ 	mpi.world.barrier();	
--~ 	PSO.print_swarm_array(PSO.swarm)
	mpi.world.barrier();

	----------------------------------------------------------------------------------------------------------------------------------------
        -- evaluating position & quality ---
	local reqs = mpi.request_list:new();
	local vals = {};
	for jj = 1 , self.swarm_size do
		if ( jj % mpi.world.size() ) == mpi.world.rank()  then
			 vals[jj] = std.vector:new( 1 );
			 local val;
			 print( "num_iter:",self.num_iter)
			 local cache_res = self:find_cache( self.swarm[jj] );			 
		         if cache_res ~= nil then
			     val = cache_res[ #cache_res ];
		        else
  			     val  = self:envaluation(self.swarm[jj]);
			end
			vals[jj][0] = val;	
			if mpi.world.rank() ~= 0 then			
				req = mpi.world.isend(0,jj,vals[jj]);
				reqs:push_back( req );		
--~ 				print( mpi.world.rank(), "send to", 0, "tag:" , jj ,val)				
			end
			
		end
		if(  mpi.world.rank()==0 ) then
			if( jj % mpi.world.size() ~= 0 ) then
				vals[jj] = std.vector:new();
				req = mpi.world.irecv( jj % mpi.world.size(),  jj, vals[jj] );
				reqs:push_back( req );	
--~ 				print( mpi.world.rank(), "recv from", jj % mpi.world.size(), "tag:", jj  )	
			end
		end
	end
	----------------------------------------------------------------------------------------------------------------------------------------
	
	mpi.wait_all(reqs);
	mpi.world.barrier();
	
	
	
	
	----------------------------------------------------------------------------------------------------------------------------------------
	--broadcast values
	for jj = 1 , self.swarm_size do	  
	   if mpi.world.rank() == 0 then	
		mpi.world.broadcast(vals[jj],0);
	   else
	        local val = std.vector:new();
		mpi.world.broadcast(val,0);
		vals[jj] = val;
	   end
	end
	----------------------------------------------------------------------------------------------------------------------------------------
	
	----------------------------------------------------------------------------------------------------------------------------------------
	-- update val and save array
	for jj = 1 , self.swarm_size do
		self.swarm[jj].value = vals[jj][0]
		if mpi.world.rank() == 0 then	
			self:write2file( self.swarm[jj] );
		end
	end
		
	----------------------------------------------------------------------------------------------------------------------------------------
	-- update best position and val
	for jj = 1 , self.swarm_size do			
		if vals[jj][0] < self.swarm[jj].best_value then              -- if new position is better
		    for k,v in pairs(self.par) do
			 self.swarm[jj][k].best_pos =  self.swarm[jj][k].pos ; -- update best position
		    end
		    self.swarm[jj].best_value= vals[jj][0];              -- % and best value
		end       
	end	    
	best_swarm, self.global_best_value= self:get_global_best_value();	
	self.best_swarm = best_swarm;	
--~ 	PSO.print_swarm_part( best_swarm);
	----------------------------------------------------------------------------------------------------------------------------------------
	
	----------------------------------------------------------------------------------------------------------------------------------------
	--- updating velocity vectors
	for jj = 1 , self.swarm_size do		
		for k,v in pairs(self.par) do
		     self.swarm[jj][k].vel = math.random() * self.inertia * self.swarm[jj][k].vel 
							    + self.correction_factor*math.random()*(self.swarm[jj][k].best_pos - self.swarm[jj][k].pos) 
							    + self.correction_factor*math.random()*(best_swarm[k].best_pos - self.swarm[jj][k].pos)
		end
	end 
	----------------------------------------------------------------------------------------------------------------------------------------	
end

function PSO:savebestinfo(path)
	
	local filename = path.."/best"..self.num_iter..".txt";
	lfs.mkdir(path.."/");
	lfs.touch( filename );
	local best_thisround = assert( io.open( filename,"w"))
	best_thisround:write("��Ӧ�ȣ�"..self.best_swarm.best_value.."\n")
	for k,v in pairs(self.par) do
		best_thisround:write(v.name, string.format("=%20.3f", self.best_swarm[k].best_pos).."\n")		
        end
        
    
	best_thisround:close()
end

function PSO:iterate()
	for ii=1,self.iterations do	    
		self:step();
                self:savebestinfo("data/best");
	end
end

function PSO:run()
    opt:read_cache( "cache.txt" );
    opt:init();
    PSO:iterate();
end



